"""
# @Time    : 2021/7/1 8:44 上午
# @Author  : hezhiqiang01
# @Email   : hezhiqiang01@baidu.com
# @File    : env_wrappers.py
从OpenAI Baselines代码修改，使其适用于多智能体环境
"""

import numpy as np

# 单个环境
class DummyVecEnv():
    """
    虚拟向量化环境类，用于并行运行多个环境实例
    """
    def __init__(self, env_fns):
        """
        初始化虚拟向量化环境
        
        参数:
            env_fns: 环境创建函数的列表
        """
        self.envs = [fn() for fn in env_fns]  # 创建环境实例列表
        env = self.envs[0]  # 获取第一个环境实例作为参考
        self.num_envs = len(env_fns)  # 环境数量
        self.observation_space = env.observation_space  # 观察空间
        self.share_observation_space = env.share_observation_space  # 共享观察空间
        self.action_space = env.action_space  # 动作空间
        self.actions = None  # 存储动作

    def step(self, actions):
        """
        同步执行环境的步进。
        这是为了向后兼容而提供的。
        
        参数:
            actions: 要执行的动作
            
        返回:
            观察、奖励、完成标志和信息
        """
        self.step_async(actions)  # 异步执行步进
        return self.step_wait()  # 等待步进完成并返回结果

    def step_async(self, actions):
        """
        异步执行环境的步进
        
        参数:
            actions: 要执行的动作
        """
        self.actions = actions  # 存储动作

    def step_wait(self):
        """
        等待异步步进完成并返回结果
        
        返回:
            观察、奖励、完成标志和信息
        """
        # 对每个环境执行步进
        results = [env.step(a) for (a, env) in zip(self.actions, self.envs)]
        obs, rews, dones, infos = map(np.array, zip(*results))  # 解包结果

        # 对于已完成的环境，重置环境
        for (i, done) in enumerate(dones):
            if 'bool' in done.__class__.__name__:  # 如果done是布尔类型
                if done:
                    obs[i] = self.envs[i].reset()  # 重置环境
            else:  # 如果done是数组
                if np.all(done):  # 如果所有元素都为True
                    obs[i] = self.envs[i].reset()  # 重置环境

        self.actions = None  # 清空动作
        return obs, rews, dones, infos  # 返回结果

    def reset(self):
        """
        重置所有环境
        
        返回:
            初始观察
        """
        obs = [env.reset() for env in self.envs]  # [env_num, agent_num, obs_dim]
        return np.array(obs)  # 转换为numpy数组

    def close(self):
        """
        关闭所有环境
        """
        for env in self.envs:
            env.close()  # 关闭环境

    def render(self, mode="human"):
        """
        渲染环境
        
        参数:
            mode: 渲染模式，可以是"human"或"rgb_array"
            
        返回:
            如果mode是"rgb_array"，返回渲染的图像数组
        """
        if mode == "rgb_array":
            # 返回所有环境的渲染图像数组
            return np.array([env.render(mode=mode) for env in self.envs])
        elif mode == "human":
            # 以人类可视化方式渲染所有环境
            for env in self.envs:
                env.render(mode=mode)
        else:
            # 不支持的渲染模式
            raise NotImplementedError
